import gradio as gr
import pickle
import numpy as np
from PIL import Image, ImageOps

# 1. 加载预先训练好的最佳 KNN 模型
try:
    with open('best_knn_model.pkl', 'rb') as f:
        knn_model = pickle.load(f)
    print("[模型加载成功] best_knn_model.pkl")
except FileNotFoundError:
    print("错误：'best_knn_model.pkl' 模型文件未找到。")
    print("请先运行 'optimal_knn.py' 脚本来训练并生成模型文件。")
    exit()

# 2. 定义核心的预测函数
def predict_digit(raw_input):
    """
    接收来自Gradio的任何格式的输入，进行净化、预处理和预测。
    """
    if raw_input is None:
        return "请先在画板上绘制一个数字"
    
    # ####################################################################
    # ## 根本性修正：输入净化器 ##
    # ####################################################################
    # 无论输入是什么格式，都将其转换为 PIL 图像对象
    
    pil_image = None
    
    # 情况1: 输入是字典 (最常见的问题)
    if isinstance(raw_input, dict):
        # 尝试从常见的键 'image' 或 'composite' 中提取图像数据
        if 'image' in raw_input and raw_input['image'] is not None:
            pil_image = Image.fromarray(raw_input['image'])
        elif 'composite' in raw_input and raw_input['composite'] is not None:
            pil_image = Image.fromarray(raw_input['composite'])
        else:
            return "错误：无法从Gradio的字典输出中找到图像数据"
            
    # 情况2: 输入是 NumPy 数组
    elif isinstance(raw_input, np.ndarray):
        pil_image = Image.fromarray(raw_input)
        
    # 情况3: 输入已经是 PIL 图像对象
    elif isinstance(raw_input, Image.Image):
        pil_image = raw_input
        
    else:
        return f"错误：收到了无法识别的输入类型: {type(raw_input)}"
    # ####################################################################
    
    # --- 从这里开始，pil_image 一定是一个有效的 PIL 图像对象 ---

    # 步骤 1: 转换为灰度图
    gray_image = pil_image.convert('L')

    # 步骤 2: 颜色反转
    inverted_image = ImageOps.invert(gray_image)

    # 步骤 3: 裁剪边界框
    bbox = inverted_image.getbbox()
    if bbox is None:
        return "无法检测到笔迹，请重新绘制"
    cropped_image = inverted_image.crop(bbox)
    
    # 步骤 4: 保持比例，粘贴到正方形背景中
    width, height = cropped_image.size
    max_dim = max(width, height)
    square_image = Image.new('L', (max_dim, max_dim), 0)
    paste_x = (max_dim - width) // 2
    paste_y = (max_dim - height) // 2
    square_image.paste(cropped_image, (paste_x, paste_y))

    # 步骤 5: 缩放并转换为模型需要的格式
    resized_image = square_image.resize((8, 8), Image.Resampling.LANCZOS)
    final_array = np.array(resized_image)
    scaled_array = final_array / 255.0 * 16.0
    flattened_array = scaled_array.flatten()
    input_data = flattened_array.reshape(1, -1)

    # --- 模型预测 ---
    prediction = knn_model.predict(input_data)
    
    return int(prediction[0])

# 3. 创建并配置Gradio网页应用
webapp = gr.Interface(
    fn=predict_digit,
    # 我们将 type 设置回 'numpy'，因为我们的净化器可以处理它可能返回的字典
    inputs=gr.Sketchpad(label="手写输入区", type="numpy"),
    outputs=gr.Label(num_top_classes=1, label="预测结果"),
    title="手写数字识别系统 (KNN最优模型)",
    description="在左侧画板中手写一个数字 (0-9)，然后点击“提交”按钮。系统将显示识别结果。",
    live=False,
    submit_btn=gr.Button("提交", variant="primary"),
    clear_btn=None
)

# 4. 启动应用
if __name__ == "__main__":
    print("[系统启动] 请在浏览器打开显示的URL进行手写测试。")
    webapp.launch()